from fastapi import FastAPI, Request, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
import asyncio
import json
import time
import uuid
import threading
from typing import Dict, Any, Optional, List

# 先初始化基本组件，避免事件循环冲突
from utils.logger import Logger
from utils.constants import ErrorCode

# 初始化日志
logger = Logger.get_logger(agent_name="text_qa_agent")

# 创建FastAPI应用
app = FastAPI(
    title="Text QA Agent API",
    description="文本问答智能体服务接口",
    docs_url="/docs"
)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境应该设置具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Agent实例和相关组件（延迟初始化）
text_qa_agent = None
async_text_qa_agent = None
redis_client = None
import uuid
from communication.agent_registry import agent_registry, AgentInfo, AgentCapability, AgentStatus
# 延迟导入AgentState以避免事件循环问题
from datetime import datetime

agent_id = f"text_qa_{uuid.uuid4().hex[:8]}"
agent_type = "text_qa"  # 确保与编排器期望的类型匹配
agent_name = "文本问答智能体"
agent_version = "1.0.0"
agent_endpoint = "http://localhost:8001"

# Agent状态对象，全局变量
agent_state = None

# 使用AgentState类注册Agent
async def register_to_orchestrator():
    """使用AgentState类注册Agent到Redis，确保被编排器正确发现"""
    print("\n--------------------------------------------------------")
    print("                     REGISTRATION START                 ")
    print("--------------------------------------------------------")
    
    global agent_state
    try:
        logger.info(f"开始使用AgentState类注册Agent: {agent_id} | {agent_type} | {agent_endpoint}")
        print(f"[REGISTER] 开始注册Agent: {agent_id} | {agent_type} | {agent_endpoint}")
        
        # 动态导入AgentState以避免事件循环问题
        print("[REGISTER] 导入AgentState类...")
        from states.agent_state import AgentState, AgentHealthStatus
        print("[REGISTER] AgentState类导入成功")
        
        # 额外日志：验证AgentType
        logger.info(f"验证AgentType: {agent_type}")
        print(f"[REGISTER] 验证AgentType: {agent_type}")
        
        # 创建AgentState实例
        logger.info(f"正在创建AgentState实例...")
        print(f"[REGISTER] 创建AgentState实例: agent_type={agent_type}, instance_id={agent_id}, port=8001")
        # 首先声明global
        global agent_state
        # 然后赋值
        agent_state = AgentState(
            agent_type=agent_type,
            instance_id=agent_id,
            port=8001  # 端口号
        )
        
        redis_key = getattr(agent_state, 'redis_key', '未知')
        logger.info(f"AgentState实例创建成功，Redis键: {redis_key}")
        print(f"[REGISTER] AgentState实例创建成功，Redis键: {redis_key}")
        print(f"[REGISTER] AgentState属性: {vars(agent_state)}")
        
        # 强制设置为健康状态
        print("[REGISTER] 设置Agent为健康状态...")
        agent_state.status = AgentHealthStatus.HEALTHY.value
        agent_state.fail_count = 0
        agent_state.task_success_rate = 100.0
        logger.info(f"Agent状态已设置: {agent_state.status}")
        print(f"[REGISTER] Agent状态已设置: {agent_state.status}")
        
        # 立即同步到Redis
        logger.info(f"正在同步到Redis...")
        print("[REGISTER] 同步状态到Redis...")
        try:
            success = agent_state._sync_to_redis(register=True)
            logger.info(f"AgentState注册结果: {success} | agent_id={agent_id} | agent_type={agent_type}")
            print(f"[REGISTER] Redis同步结果: {success}")
        except Exception as sync_error:
            logger.error(f"Redis同步失败: {str(sync_error)}")
            print(f"[REGISTER] ERROR: Redis同步失败: {str(sync_error)}")
            success = False
        
        # 手动发送一次心跳确保状态正确
        logger.info(f"手动发送心跳...")
        print("[REGISTER] 发送初始化心跳...")
        try:
            heartbeat_success = agent_state.send_heartbeat()
            logger.info(f"发送初始化心跳: {heartbeat_success}")
            print(f"[REGISTER] 心跳发送结果: {heartbeat_success}")
        except Exception as heartbeat_error:
            logger.error(f"心跳发送失败: {str(heartbeat_error)}")
            print(f"[REGISTER] ERROR: 心跳发送失败: {str(heartbeat_error)}")
        
        # 添加Redis直接检查
        logger.info(f"直接检查Redis状态...")
        print("[REGISTER] 直接检查Redis状态...")
        try:
            from middleware.redis_adapter import get_redis_adapter
            redis_client = get_redis_adapter()
            print(f"[REGISTER] 获取Redis客户端: {redis_client}")
            
            # 直接从Redis读取状态验证注册是否成功
            if redis_client:
                # 移除await，Redis客户端可能是同步的
                redis_data = redis_client.hgetall(redis_key)
                found = len(redis_data) > 0
                logger.info(f"Redis读取结果: 找到数据={found} | 数据预览={str(redis_data)[:200]}")
                print(f"[REGISTER] Redis检查结果: 找到数据={found}")
                if found:
                    print(f"[REGISTER] Redis数据示例: {dict(redis_data) if len(redis_data) < 10 else '数据量较大'}")
                else:
                    print(f"[REGISTER] WARNING: Redis中没有找到Agent数据!")
            else:
                print("[REGISTER] ERROR: Redis客户端不可用")
        except Exception as redis_e:
            logger.error(f"Redis检查失败: {str(redis_e)}")
            print(f"[REGISTER] ERROR: Redis检查失败: {str(redis_e)}")
        
        # 额外检查agent_health键
        try:
            health_key = f"agent_health:{agent_type}"
            health_data = redis_client.get(health_key) if redis_client else None
            print(f"[REGISTER] 检查健康状态键 {health_key}: {health_data}")
        except Exception as e:
            print(f"[REGISTER] 健康状态检查失败: {str(e)}")
        
        print(f"[REGISTER] 注册完成，结果: {success}")
        return success
    except Exception as e:
        logger.error(f"AgentState注册失败: {str(e)}", exc_info=True)
        print(f"[REGISTER] ERROR: 注册失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        print("--------------------------------------------------------")
        print("                     REGISTRATION END                   ")
        print("--------------------------------------------------------\n")

# 心跳保活函数
def start_heartbeat():
    """定期发送心跳，保持Agent状态活跃"""
    def heartbeat_loop():
        while True:
            try:
                if agent_state:
                    # 直接调用同步的心跳方法
                    success = agent_state.send_heartbeat()
                    logger.debug(f"心跳发送结果: {success} | 当前状态: {agent_state.status}")
                else:
                    logger.warning("AgentState未初始化，无法发送心跳")
            except Exception as e:
                logger.error(f"心跳发送失败: {str(e)}")
            time.sleep(10)  # 每10秒发送一次心跳
    
    # 启动心跳线程
    heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True)
    heartbeat_thread.start()
    logger.info("心跳保活线程已启动")

# 启动时注册Agent
@app.on_event("startup")
async def startup_event():
    print("========================================================")
    print("                    STARTUP EVENT STARTED               ")
    print("========================================================")
    
    global text_qa_agent, async_text_qa_agent, redis_client, agent_state
    
    logger.info("Text QA Agent 服务启动中...")
    print("[STARTUP] Text QA Agent 服务启动中...")
    
    # 延迟初始化组件
    try:
        print("[STARTUP] 导入必要组件...")
        from middleware.redis_adapter import get_redis_adapter
        from apps.text_qa.text_qa_agent import TextQAAgent
        from apps.text_qa.async_text_qa_agent import AsyncTextQAAgent
        print("[STARTUP] 组件导入成功")
        
        # 首先初始化云端LLM模型客户端
        logger.info("开始初始化云端LLM模型客户端...")
        try:
            from middleware.model_client import initialize_model_clients
            # 确保使用await而不是asyncio.run
            await initialize_model_clients()
            logger.info("云端LLM模型客户端初始化完成")
        except Exception as e:
            logger.error(f"云端LLM模型客户端初始化失败: {str(e)}")
            # 继续执行，尝试使用备用方案
            
        # 确保模型适配器初始化
        logger.info("开始初始化模型适配器...")
        try:
            from middleware.model_adapter import model_adapter
            # 如果模型适配器需要初始化，确保它被正确初始化
            if hasattr(model_adapter, 'initialize'):
                if asyncio.iscoroutinefunction(model_adapter.initialize):
                    await model_adapter.initialize()
                else:
                    model_adapter.initialize()
            logger.info("模型适配器初始化完成")
        except Exception as e:
            logger.error(f"模型适配器初始化失败: {str(e)}")
            # 继续执行，尝试使用备用方案
        
        print("[STARTUP] 初始化TextQAAgent...")
        text_qa_agent = TextQAAgent()
        print("[STARTUP] 初始化AsyncTextQAAgent...")
        async_text_qa_agent = AsyncTextQAAgent()
        print("[STARTUP] 初始化Redis客户端...")
        redis_client = await get_redis_adapter() if hasattr(get_redis_adapter, '__call__') and asyncio.iscoroutinefunction(get_redis_adapter) else get_redis_adapter()
        print(f"[STARTUP] Redis客户端初始化完成: {redis_client}")
        
        logger.info(f"Text QA Agent 初始化成功: agent_id={agent_id}")
        print(f"[STARTUP] Text QA Agent 初始化成功: agent_id={agent_id}")
        
        # 使用AgentState注册到Redis
        print("[STARTUP] 开始注册Agent到Redis...")
        if await register_to_orchestrator():
            # 启动心跳保活线程
            print("[STARTUP] 注册成功，启动心跳保活线程...")
            start_heartbeat()
            logger.info("Text QA Agent 服务启动成功，心跳保活已启动")
            print("[STARTUP] Text QA Agent 服务启动成功，心跳保活已启动")
        else:
            logger.error("Text QA Agent 服务启动失败：注册失败")
            print("[STARTUP] ERROR: Text QA Agent 服务启动失败：注册失败")
    except Exception as e:
        logger.error(f"Text QA Agent 初始化失败: {str(e)}", exc_info=True)
        print(f"[STARTUP] ERROR: Text QA Agent 初始化失败: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        print("========================================================")
        print("                    STARTUP EVENT COMPLETED            ")
        print("========================================================")

# 关闭时注销Agent
@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Text QA Agent 服务关闭中...")
    global agent_state
    try:
        # 使用AgentState注销
        if agent_state:
            # 确保有AgentState导入
            from states.agent_state import AgentState, AgentHealthStatus
            agent_state.unregister()
            logger.info(f"Text QA Agent 已从Redis注销: {agent_id}")
        
        # 同时使用标准注册中心注销
        success = agent_registry.unregister_agent(agent_id)
        if success:
            logger.info(f"Text QA Agent 已从标准注册中心注销: {agent_id}")
        else:
            logger.warning(f"Text QA Agent 从标准注册中心注销失败: {agent_id} (可能已不存在)")
    except Exception as e:
        logger.error(f"Text QA Agent 注销失败: {str(e)}")

# 健康检查接口
@app.get("/", tags=["基础检查"])
async def root():
    return {
        "code": ErrorCode.SUCCESS.value,
        "message": "Text QA Agent 服务运行中",
        "data": {
            "agent_id": agent_id,
            "agent_name": agent_name,
            "agent_type": agent_type,
            "timestamp": time.time()
        }
    }

# 预测接口（处理任务）
@app.post("/predict", tags=["核心功能"])
async def predict(request: Request, background_tasks: BackgroundTasks) -> Dict[str, Any]:
    """处理预测请求"""
    global text_qa_agent, async_text_qa_agent
    
    try:
        # 确保Agent已初始化
        if text_qa_agent is None or async_text_qa_agent is None:
            # 报告真实的初始化错误，而不是使用模拟对象
            error_msg = "Text QA Agent 初始化失败，请检查系统日志获取详细错误信息"
            logger.error(error_msg)
            return {
                "code": ErrorCode.SYSTEM_ERROR.value,
                "message": error_msg,
                "data": None
            }
        
        # 暂时移除模型适配器的直接测试，避免初始化问题
        # 模型适配器的实际使用会在任务处理中进行，到时候再捕获具体错误
        logger.info("跳过模型适配器预测试，直接处理请求")
        
        # 获取请求数据
        task = await request.json()
        task_id = task.get("sub_task_id", "unknown")
        query = task.get("query", "")
        
        # 更新日志上下文
        Logger.update_context(task_id=task_id, agent_name=agent_name)
        logger.info(f"收到预测请求: query={query[:50]}...")
        
        # 异步处理任务
        try:
            print(f"[CRITICAL DEBUG] 准备调用async_text_qa_agent.process_task_async，任务数据: {task}")
            # 使用异步Agent处理任务
            result = await async_text_qa_agent.process_task_async(task)
            print(f"[CRITICAL DEBUG] async_text_qa_agent.process_task_async调用成功，结果: {result}")
            
            logger.info(f"任务处理成功: task_id={task_id}")
            return {
                "code": ErrorCode.SUCCESS.value,
                "message": "任务处理成功",
                "data": result
            }
        except Exception as e:
            error_type = type(e).__name__
            error_message = str(e)
            print(f"[CRITICAL DEBUG] 异步处理失败: {error_type}: {error_message}")
            import traceback
            print(f"[CRITICAL DEBUG] 异步处理错误堆栈:\n{traceback.format_exc()}")
            # 如果异步处理失败，尝试同步处理
            logger.warning(f"异步处理失败，尝试同步处理: {error_type}: {error_message}")
            try:
                print(f"[CRITICAL DEBUG] 尝试同步处理任务")
                result = text_qa_agent.process_task(task)
                print(f"[CRITICAL DEBUG] 同步处理成功，结果: {result}")
                logger.info(f"同步任务处理成功: task_id={task_id}")
                return {
                    "code": ErrorCode.SUCCESS.value,
                    "message": "任务处理成功",
                    "data": result
                }
            except Exception as sync_e:
                sync_error_type = type(sync_e).__name__
                sync_error_message = str(sync_e)
                print(f"[CRITICAL DEBUG] 同步处理也失败: {sync_error_type}: {sync_error_message}")
                print(f"[CRITICAL DEBUG] 同步处理错误堆栈:\n{traceback.format_exc()}")
                raise
    
    except Exception as e:
        logger.error(f"任务处理失败: {str(e)}", exc_info=True)
        return {
            "code": ErrorCode.SYSTEM_ERROR.value,
            "message": f"任务处理失败: {str(e)}",
            "data": None
        }

# 任务处理接口（兼容旧版本）
@app.post("/process", tags=["核心功能"])
async def process(request: Request) -> Dict[str, Any]:
    """处理任务请求（兼容接口）"""
    return await predict(request, None)

# 编排器调用接口（/api/run）
@app.post("/api/run", tags=["编排器接口"])
async def api_run(request: Request, background_tasks: BackgroundTasks) -> Dict[str, Any]:
    """处理编排器的任务请求（AsyncOrchestratorAgent调用）"""
    global text_qa_agent, async_text_qa_agent
    
    try:
        # 获取请求数据
        task = await request.json()
        
        # 确保任务有task_id
        if "task_id" not in task:
            task["task_id"] = str(uuid.uuid4())
        task_id = task["task_id"]
        
        # 更新日志上下文
        Logger.update_context(task_id=task_id, agent_name=agent_name)
        logger.info(f"收到编排器任务请求: task_id={task_id}")
        
        # 适配异步编排器的请求格式
        # 优先从inputs字段中获取query，同时兼容直接query和input字段
        query_text = task.get("inputs", {}).get("query", "")
        if not query_text:
            query_text = task.get("query", "")
        if not query_text:
            query_text = task.get("input", "")
        
        # 确保任务中有query字段
        task["query"] = query_text
        
        # 确保任务中有基本字段
        if "agent_type" not in task:
            task["agent_type"] = "text_qa"
        if "timestamp" not in task:
            task["timestamp"] = time.time()
        
        # 记录处理任务的调试信息
        logger.info(f"[DEBUG] 处理编排器任务: task_id={task_id}, query={query_text[:50]}...")
        
        print(f"[CRITICAL DEBUG] 准备调用async_text_qa_agent.process_task_async处理编排器任务，完整任务数据: {task}")
        # 使用异步Agent处理完整的任务对象
        result = await async_text_qa_agent.process_task_async(task)
        print(f"[CRITICAL DEBUG] 编排器任务处理完成，原始结果: {result}")
        
        # 确保返回格式符合预期
        if "status" not in result and "success" not in result:
            result["success"] = True
        
        logger.info(f"编排器任务处理成功: task_id={task_id}")
        return result
    except Exception as e:
        logger.error(f"编排器任务处理失败: {str(e)}", exc_info=True)
        # 确保返回编排器期望的格式，使用正确的错误代码
        return {
            "success": False,
            "error": str(e),
            "code": ErrorCode.AGENT_CALL_ERROR.value
        }

# 启动入口
if __name__ == "__main__":
    logger.info(f"启动Text QA Agent服务，监听端口: 8001")
    uvicorn.run(
        "apps.text_qa.app:app",
        host="0.0.0.0",
        port=8001,
        reload=True,
        log_config=None
    )